#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelGelu
{
public:
    __aicore__ inline KernelGelu() {}
    __aicore__ void Init(GM_ADDR x, GM_ADDR y,
                         uint32_t blockLength, uint32_t totalLength,
                         uint32_t tileNum, uint32_t tileLength,
                         uint32_t lasttileLength)
    {
        this->blockLength = blockLength;
        this->lasttileLength = lasttileLength;
        this->tileLength = tileLength / BUFFER_NUM;
        this->tileNum = tileNum;

        uint32_t offset = this->blockLength * GetBlockIdx();
        xGm.SetGlobalBuffer((__gm__ DTYPE_X *)x + offset,
                            this->blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y *)y + offset,
                            this->blockLength);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(inQueueX0, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(inQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(tmp1, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(DTYPE_Y));
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++)
        {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> x0Local = inQueueX0.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> x1Local = inQueueX1.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> x2Local = inQueueX1.AllocTensor<DTYPE_X>();
        LocalTensor<DTYPE_X> y0Local = inQueueY.AllocTensor<DTYPE_X>();

        if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
            (progress == (this->tileNum * BUFFER_NUM - 1)))
        {
            DataCopy(xLocal, xGm[(progress - 2) * (this->tileLength) + this->lasttileLength], this->tileLength);
            DataCopy(x0Local, xGm[(progress - 2) * (this->tileLength) + this->lasttileLength], this->tileLength);
            DataCopy(x1Local, xGm[(progress - 2) * (this->tileLength) + this->lasttileLength], this->tileLength);
            DataCopy(x2Local, xGm[(progress - 2) * (this->tileLength) + this->lasttileLength], this->tileLength);
            DataCopy(y0Local, xGm[(progress - 2) * (this->tileLength) + this->lasttileLength], this->tileLength);
        }
        else
        {
            DataCopy(xLocal, xGm[progress * (this->tileLength)], this->tileLength);
            DataCopy(x0Local, xGm[progress * (this->tileLength)], this->tileLength);
            DataCopy(x1Local, xGm[progress * (this->tileLength)], this->tileLength);
            DataCopy(x2Local, xGm[progress * (this->tileLength)], this->tileLength);
            DataCopy(y0Local, xGm[progress * (this->tileLength)], this->tileLength);
        }

        inQueueX.EnQue(xLocal);
        inQueueX0.EnQue(x0Local);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
        inQueueY.EnQue(y0Local);
    } // CopyIn
    __aicore__ inline void Compute(int32_t progress)
    {

        if constexpr (std::is_same_v<DTYPE_X, float>)
        {
            LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> x0Local = inQueueX0.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> x1Local = inQueueX1.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> x2Local = inQueueX2.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> y0Local = inQueueY.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> yLocal = outQueueY.AllocTensor<DTYPE_X>();
            DTYPE_X p = 0.3275911;
            DTYPE_X a1 = 0.254829592;
            DTYPE_X a2 = -0.284496736;
            DTYPE_X a3 = 1.421413741;
            DTYPE_X a4 = -1.453152027;
            DTYPE_X a5 = 1.061405429;
            Abs(xLocal, xLocal, this->tileLength);
            Div(x0Local, x0Local, xLocal, this->tileLength);             // x0 = sign(x)
            Muls(xLocal, xLocal, (DTYPE_X)0.70710677, this->tileLength); // x = abs(x)*0.70710677
            Muls(yLocal, xLocal, p, this->tileLength);
            Adds(yLocal, yLocal, (DTYPE_X)1, this->tileLength);
            Reciprocal(xLocal, yLocal, this->tileLength); // x = t
            Muls(yLocal, xLocal, a1, this->tileLength);
            Mul(x1Local, xLocal, xLocal, this->tileLength);
            Muls(x2Local, x1Local, a2, this->tileLength);
            Add(yLocal, yLocal, x2Local, this->tileLength);
            Mul(x1Local, x1Local, xLocal, this->tileLength);
            Muls(x2Local, x1Local, a3, this->tileLength);
            Add(yLocal, yLocal, x2Local, this->tileLength);
            Mul(x1Local, x1Local, xLocal, this->tileLength);
            Muls(x2Local, x1Local, a4, this->tileLength);
            Add(yLocal, yLocal, x2Local, this->tileLength);
            Mul(x1Local, x1Local, xLocal, this->tileLength);
            Muls(x1Local, x1Local, a5, this->tileLength);
            Add(yLocal, yLocal, x1Local, this->tileLength);
            Muls(x1Local, x1Local, (DTYPE_X)0, this->tileLength);
            Mul(xLocal, y0Local, y0Local, this->tileLength);
            Sub(xLocal, x1Local, xLocal, this->tileLength);
            Exp(xLocal, xLocal, this->tileLength);
            Mul(yLocal, yLocal, xLocal, this->tileLength);
            Adds(x1Local, x1Local, (DTYPE_X)1, this->tileLength);
            Sub(yLocal, x1Local, yLocal, this->tileLength);
            Mul(yLocal, yLocal, x0Local, this->tileLength);
            Adds(yLocal, yLocal, (DTYPE_X)1, this->tileLength);
            Mul(yLocal, yLocal, y0Local, this->tileLength);
            Muls(yLocal, yLocal, (DTYPE_X)0.5, this->tileLength);

            // DTYPE_X a0 = 0.1400122886866665;
            // DTYPE_X a1 = 1.2732395447351628;
            // Muls(xLocal, xLocal, (DTYPE_X)0.70710677, this->tileLength);
            // Mul(xLocal, xLocal, xLocal, this->tileLength);
            // Muls(xLocal, xLocal, a0, this->tileLength);
            // Adds(yLocal, xLocal, a1, this->tileLength);
            // Adds(xLocal, xLocal, (DTYPE_X)1, this->tileLength);
            // Div(yLocal, yLocal, xLocal, this->tileLength);
            // Mul(yLocal, yLocal, x0Local, this->tileLength);
            // Mul(yLocal, yLocal, x0Local, this->tileLength);
            // Muls(yLocal, yLocal, (DTYPE_X)0.70710677, this->tileLength);
            // Muls(yLocal, yLocal, (DTYPE_X)0.70710677, this->tileLength);
            // Muls(xLocal, xLocal, (DTYPE_X)0, this->tileLength);
            // Sub(yLocal, xLocal, yLocal, this->tileLength);
            // Exp(yLocal, yLocal, this->tileLength);
            // Adds(xLocal, xLocal, (DTYPE_X)1, this->tileLength);
            // Sub(yLocal, xLocal, yLocal, this->tileLength);
            // Sqrt(yLocal, yLocal, this->tileLength);
            // Abs(xLocal, x0Local, this->tileLength);
            // Div(xLocal, x0Local, xLocal, this->tileLength);
            // Mul(yLocal, yLocal, xLocal, this->tileLength);
            // Adds(yLocal, yLocal, (DTYPE_X)1, this->tileLength);
            // Mul(yLocal, yLocal, x0Local, this->tileLength);
            // Muls(yLocal, yLocal, (DTYPE_X)0.5, this->tileLength);
            outQueueY.EnQue<DTYPE_X>(yLocal);
            inQueueX.FreeTensor(xLocal);
            inQueueX.FreeTensor(x0Local);
            inQueueX.FreeTensor(x1Local);
            inQueueX.FreeTensor(x2Local);
            inQueueX.FreeTensor(y0Local);
        }
        else
        {
            auto tmp1Local = tmp1.Get<float>();
            auto tmp2Local = tmp2.Get<float>();
            LocalTensor<DTYPE_X> xLocal = inQueueX.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> x0Local = inQueueX0.DeQue<DTYPE_X>();
           LocalTensor<DTYPE_X> x1Local = inQueueX1.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> x2Local = inQueueX2.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_X> y0Local = inQueueY.DeQue<DTYPE_X>();
            LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
            Cast(tmp1Local, xLocal, RoundMode::CAST_NONE, this->tileLength);
            Mul(tmp2Local, tmp1Local, tmp1Local, this->tileLength);
            Mul(tmp2Local, tmp2Local, tmp1Local, this->tileLength);
            Muls(tmp2Local, tmp2Local, (float)0.044715, this->tileLength);
            Add(tmp2Local, tmp2Local, tmp1Local, this->tileLength);
            Muls(tmp2Local, tmp2Local, (float)(-1.59576912), this->tileLength); // -1.59576912(x+0.044715x^3)
            Exp(tmp2Local, tmp2Local, this->tileLength);
            Adds(tmp2Local, tmp2Local, (float)(1), this->tileLength);
            Div(tmp2Local, tmp1Local, tmp2Local, this->tileLength);
            Cast(yLocal, tmp2Local, RoundMode::CAST_NONE, this->tileLength);
            outQueueY.EnQue<DTYPE_Y>(yLocal);
            inQueueX.FreeTensor(xLocal);
            inQueueX.FreeTensor(x0Local);
            inQueueX.FreeTensor(x1Local);
            inQueueX.FreeTensor(x2Local);
            inQueueX.FreeTensor(y0Local);
        }
    } // Compute

    __aicore__ inline void CopyOut(int32_t progress)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
            (progress == (this->tileNum * BUFFER_NUM - 1)))
        {
            DataCopy(yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                     yLocal, (this->tileLength));
        }
        else
        {
            DataCopy(yGm[progress * (this->tileLength)], yLocal,
                     (this->tileLength));
        }
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX, inQueueX0, inQueueX1,inQueueX2, inQueueY;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> tmp1, tmp2;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;
};

extern "C" __global__ __aicore__ void gelu(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelGelu op;
    op.Init(x, y, tiling_data.blockLength,
            tiling_data.totalLength, tiling_data.tileNum,
            tiling_data.tileLength, tiling_data.lasttileLength);
    op.Process();
}
 
